{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pretty printing has been turned OFF\n",
      "Pretty printing has been turned ON\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pprint\n",
    "from Environment import GridworldEnv\n",
    "from pprint import PrettyPrinter\n",
    "\n",
    "%pprint\n",
    "pp = PrettyPrinter(indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "env = GridworldEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def policy_eval(policy, environment, discount_factor=1.0, theta=0.1):\n",
    "    env = environment # 环境变量\n",
    "    \n",
    "    # 初始化一个全0的价值函数\n",
    "    V = np.zeros(env.nS)\n",
    "    \n",
    "    # 迭代开始\n",
    "    for _ in range(50):\n",
    "        delta = 0\n",
    "        \n",
    "        # 对于GridWorld中的每一个状态都进行全备份\n",
    "        for s in range(env.nS):\n",
    "            v = 0\n",
    "            # 检查下一个有可能执行的动作\n",
    "            for a, action_prob in enumerate(policy[s]):\n",
    "                \n",
    "                # 对于每一个动作检查下一个状态\n",
    "                for  prob, next_state, reward, done in env.P[s][a]:\n",
    "                    # 累积计算下一个动作的期望价值\n",
    "                    v += action_prob * prob * (reward + discount_factor * V[next_state])\n",
    "            # 选出最大的变化量\n",
    "            delta = max(delta, np.abs(v - V[s]))\n",
    "            V[s] = v\n",
    "        \n",
    "        # 停止标志位\n",
    "        if delta <= theta:\n",
    "            break\n",
    "    \n",
    "    return np.array(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def policy_improvement(env, policy, discount_factor=1.0):\n",
    "    \"\"\"\n",
    "    Policy Imrpovement.\n",
    "    Iterativedly evaluates and improves a policy until an \n",
    "    optimal policy is found or to the limited iter threshold.\n",
    "    \n",
    "    Args:\n",
    "        env: the environment.\n",
    "        policy_eval_fun: Policy Evaluation function with 3 \n",
    "        argements: policy, env, discount_factor.\n",
    "        \n",
    "    Returns:\n",
    "        tuple(policy, V).\n",
    "    \"\"\"\n",
    "    k = 0\n",
    "    while True:\n",
    "        print(k)\n",
    "        V = policy_eval(policy, env, discount_factor)\n",
    "        print(\"random policy:\\n\", policy)\n",
    "        print(\"policy eval:\\n\",V.reshape(env.shape))\n",
    "        policy_stable = True\n",
    "        for s in range(env.nS):\n",
    "            chosen_a = np.argmax(policy[s])\n",
    "            \n",
    "            action_values = np.zeros(env.nA)\n",
    "            for a in range(env.nA):\n",
    "                for prob, next_state, reward, done in env.P[s][a]:\n",
    "                    action_values[a] += prob * (reward + discount_factor * V[next_state])\n",
    "                    if done and next_state != 15:\n",
    "                        action_values[a] = float('-inf')\n",
    "\n",
    "            print(\"action_values:\\n\",s, action_values)\n",
    "            \n",
    "            best_a = np.argmax(action_values)\n",
    "            \n",
    "            if chosen_a != best_a:\n",
    "                policy_stable = False\n",
    "            policy[s] = np.eye(env.nA)[best_a]\n",
    "        \n",
    "        print(\"policy\\n\", np.reshape(np.argmax(policy, axis=1), env.shape))\n",
    "            \n",
    "        if policy_stable:\n",
    "            return policy, V\n",
    "        k+=1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "random policy:\n",
      " [[ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]\n",
      " [ 0.25  0.25  0.25  0.25]]\n",
      "policy eval:\n",
      " [[-50.16293984 -50.36830499 -49.87680931 -50.45935702]\n",
      " [-49.84606245 -50.         -47.65840563 -50.        ]\n",
      " [-48.27313919 -43.72740352 -38.47270705 -50.        ]\n",
      " [-50.         -35.52802106  -9.62643679  50.        ]]\n",
      "action_values:\n",
      " 0 [-51.16293984 -51.36830499 -50.84606245 -51.16293984]\n",
      "action_values:\n",
      " 1 [-51.36830499 -50.87680931         -inf -51.16293984]\n",
      "action_values:\n",
      " 2 [-50.87680931 -51.45935702 -48.65840563 -51.36830499]\n",
      "action_values:\n",
      " 3 [-51.45935702 -51.45935702         -inf -50.87680931]\n",
      "action_values:\n",
      " 4 [-51.16293984         -inf -49.27313919 -50.84606245]\n",
      "action_values:\n",
      " 5 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 6 [-50.87680931         -inf -39.47270705         -inf]\n",
      "action_values:\n",
      " 7 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 8 [-50.84606245 -44.72740352         -inf -49.27313919]\n",
      "action_values:\n",
      " 9 [        -inf -39.47270705 -36.52802106 -49.27313919]\n",
      "action_values:\n",
      " 10 [-48.65840563         -inf -10.62643679 -44.72740352]\n",
      "action_values:\n",
      " 11 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 12 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 13 [-44.72740352 -10.62643679 -36.52802106         -inf]\n",
      "action_values:\n",
      " 14 [-39.47270705  49.         -10.62643679 -36.52802106]\n",
      "action_values:\n",
      " 15 [ 51.  51.  51.  51.]\n",
      "policy\n",
      " [[2 1 2 3]\n",
      " [2 0 2 0]\n",
      " [1 2 2 0]\n",
      " [0 1 1 0]]\n",
      "1\n",
      "random policy:\n",
      " [[ 0.  0.  1.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 0.  0.  0.  1.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 1.  0.  0.  0.]]\n",
      "policy eval:\n",
      " [[ 38.  40.  42.  41.]\n",
      " [ 40. -50.  44. -50.]\n",
      " [ 42.  44.  46. -50.]\n",
      " [-50.  46.  48.  50.]]\n",
      "action_values:\n",
      " 0 [ 37.  39.  39.  37.]\n",
      "action_values:\n",
      " 1 [ 39.  41. -inf  37.]\n",
      "action_values:\n",
      " 2 [ 41.  40.  43.  39.]\n",
      "action_values:\n",
      " 3 [ 40.  40. -inf  41.]\n",
      "action_values:\n",
      " 4 [ 37. -inf  41.  39.]\n",
      "action_values:\n",
      " 5 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 6 [ 41. -inf  45. -inf]\n",
      "action_values:\n",
      " 7 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 8 [ 39.  43. -inf  41.]\n",
      "action_values:\n",
      " 9 [-inf  45.  45.  41.]\n",
      "action_values:\n",
      " 10 [ 43. -inf  47.  43.]\n",
      "action_values:\n",
      " 11 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 12 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 13 [ 43.  47.  45. -inf]\n",
      "action_values:\n",
      " 14 [ 45.  49.  47.  45.]\n",
      "action_values:\n",
      " 15 [ 51.  51.  51.  51.]\n",
      "policy\n",
      " [[1 1 2 3]\n",
      " [2 0 2 0]\n",
      " [1 1 2 0]\n",
      " [0 1 1 0]]\n",
      "2\n",
      "random policy:\n",
      " [[ 0.  1.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 0.  0.  0.  1.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  0.  1.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 1.  0.  0.  0.]]\n",
      "policy eval:\n",
      " [[ 38.  40.  42.  41.]\n",
      " [ 40. -50.  44. -50.]\n",
      " [ 42.  44.  46. -50.]\n",
      " [-50.  46.  48.  50.]]\n",
      "action_values:\n",
      " 0 [ 37.  39.  39.  37.]\n",
      "action_values:\n",
      " 1 [ 39.  41. -inf  37.]\n",
      "action_values:\n",
      " 2 [ 41.  40.  43.  39.]\n",
      "action_values:\n",
      " 3 [ 40.  40. -inf  41.]\n",
      "action_values:\n",
      " 4 [ 37. -inf  41.  39.]\n",
      "action_values:\n",
      " 5 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 6 [ 41. -inf  45. -inf]\n",
      "action_values:\n",
      " 7 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 8 [ 39.  43. -inf  41.]\n",
      "action_values:\n",
      " 9 [-inf  45.  45.  41.]\n",
      "action_values:\n",
      " 10 [ 43. -inf  47.  43.]\n",
      "action_values:\n",
      " 11 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 12 [-inf -inf -inf -inf]\n",
      "action_values:\n",
      " 13 [ 43.  47.  45. -inf]\n",
      "action_values:\n",
      " 14 [ 45.  49.  47.  45.]\n",
      "action_values:\n",
      " 15 [ 51.  51.  51.  51.]\n",
      "policy\n",
      " [[1 1 2 3]\n",
      " [2 0 2 0]\n",
      " [1 1 2 0]\n",
      " [0 1 1 0]]\n",
      "\n",
      "Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\n",
      "[[1 1 2 3]\n",
      " [2 0 2 0]\n",
      " [1 1 2 0]\n",
      " [0 1 1 0]]\n",
      "\n",
      "Reshaped Grid Value Function:\n",
      "[[ 38.  40.  42.  41.]\n",
      " [ 40. -50.  44. -50.]\n",
      " [ 42.  44.  46. -50.]\n",
      " [-50.  46.  48.  50.]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "random_policy = np.ones([env.nS, env.nA])/env.nA\n",
    "policy, v = policy_improvement(env, random_policy)\n",
    "\n",
    "print(\"\\nReshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\")\n",
    "print(np.reshape(np.argmax(policy, axis=1), env.shape))\n",
    "print(\"\")\n",
    "\n",
    "print(\"Reshaped Grid Value Function:\")\n",
    "print(v.reshape(env.shape))\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "The real policy Ieration function following\n",
    "============="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def policy_iteration(env, policy, discount_factor=1.0):\n",
    "    while True:\n",
    "        # 评估当前策略 policy\n",
    "        V = policy_eval(policy, env, discount_factor)\n",
    "\n",
    "        # policy 标志位，当某状态的策略更改后该标志位为 False\n",
    "        policy_stable = True\n",
    "        \n",
    "        # 策略改进\n",
    "        for s in range(env.nS):\n",
    "            # 在当前状态和策略下选择概率最高的动作\n",
    "            old_action = np.argmax(policy[s])\n",
    "            \n",
    "            # 在当前状态和策略下找到最优动作\n",
    "            action_values = np.zeros(env.nA)\n",
    "            for a in range(env.nA):\n",
    "                for prob, next_state, reward, done in env.P[s][a]:\n",
    "                    action_values[a] += prob * (reward + discount_factor * V[next_state])\n",
    "                    if done and next_state != 15:\n",
    "                        action_values[a] = float('-inf')\n",
    "\n",
    "            print(\"action_values:\\n\",s, action_values)\n",
    "            \n",
    "            # 采用贪婪算法更新当前策略\n",
    "            best_action = np.argmax(action_values)\n",
    "            \n",
    "            if old_action != best_action:\n",
    "                policy_stable = False\n",
    "            policy[s] = np.eye(env.nA)[best_a]\n",
    "        \n",
    "        \n",
    "        # 选择的动作不再变化，则代表策略已经稳定下来\n",
    "        if policy_stable:\n",
    "            # 返回最优策略和对应状态值\n",
    "            return policy, V"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
